#! python3
# 使用PyPDF2模块处理PDF文件

import PyPDF2
import os
from pdf2docx import Converter


def readPDF(filename):
    """
    从 PDF 提取文本
    :param filename:
    :return:
    """
    try:
        file = open(filename, 'rb')
        reader = PyPDF2.PdfFileReader(file)

        # 提取文本
        print('总页数：', reader.numPages)  # 取总页数
        print('总页数：', reader.getNumPages())  # 取总页数
        pageObj = reader.getPage(0)  # 第一页
        print("第一页文本的字符串[", pageObj.extractText(), "]")

        file.close()

    except FileNotFoundError as e:
        print("an Error : " + e)


def readPDF(infile, startIndex, endIndex, outfile):
    """
    读取PDF， 注意wb和rb
    读取infile的内容，并保存在outfile中
    :param infile:
    :param startIndex:读取起始页
    :param endIndex:读取结束页
    :param outfile:
    :return:
    """
    try:
        # 从 PDF 提取文本
        file = open(infile, 'rb')  # 二进制读取
        reader = PyPDF2.PdfReader(file)

        # 创建PDF
        pdfWriter = PyPDF2.PdfWriter()

        # 拷贝页面
        for i in range(startIndex, endIndex):
            pageObj = reader.pages[i]
            pdfWriter.add_page(pageObj)

        # 创建PDF续
        outFile = open(outfile, 'wb')  # 二进制写入
        pdfWriter.write(outFile)
        outFile.close()
        file.close()

    except FileNotFoundError as e:
        print("an Error : " + e)


def createAndCopyPDF(infile1, infile2, outfile):
    """
    创建PDF、拷贝PDF， 注意wb和rb
    将infile2的内容追加到infile1，并保存在outfile中
    :param infile1:
    :param infile2:
    :param outfile:
    :return:
    """
    try:
        # 从 PDF 提取文本
        file1 = open(infile1, 'rb')  # 二进制读取
        file2 = open(infile2, 'rb')  # 二进制读取
        reader1 = PyPDF2.PdfFileReader(file1)
        reader2 = PyPDF2.PdfFileReader(file2)

        # 创建PDF
        pdfWriter = PyPDF2.PdfFileWriter()

        # 拷贝页面
        for i in range(reader1.numPages - 3, reader1.numPages):  # 拷贝后三页
            pageObj = reader1.getPage(i)
            pdfWriter.addPage(pageObj)
        for i in range(reader2.getNumPages()):  # 拷贝全部
            pageObj = reader2.getPage(i)
            pdfWriter.addPage(pageObj)

        # 创建PDF续
        outFile = open(outfile, 'wb')  # 二进制写入
        pdfWriter.write(outFile)
        outFile.close()
        file1.close()
        file2.close()

    except FileNotFoundError as e:
        print("an Error : " + e)


def encryptOrDecryptPDF(filename, type_, pwd):
    """
    PDF解密、加密， 注意wb和rb
    :param filename: 文件名（可以带路径）
    :param type_: decrypt-解密，encrypt-加密
    :param pwd: 密码
    :return:
    """
    infile = open(filename, 'rb')  # 二进制读
    inReader = PyPDF2.PdfFileReader(infile)
    basename = os.path.basename(filename)  # 取文件名
    path = os.path.dirname(filename)  # 取路径

    # 解密
    if type_ == 'decrypt':
        if inReader.isEncrypted:  # 是否加密
            # 解密 PDF:
            # 如果PDF是有密码的，PyPDF2是支持输入密码的，但是仅限于其中两种密码算法，
            # 所以如果有密码需求的话，还需要测试下PyPDF2支不支持自己的PDF所用密码算法
            inReader.decrypt(pwd)  # pdf解密：only algorithm code 1 and 2 are supported

            # 解密后的文件名
            decryptFilename = os.path.join(path, basename[:len(basename) - 4] + "_decrypt.pdf")
            outfile = open(decryptFilename, 'wb')  # 二进制写
            outReader = PyPDF2.PdfFileWriter()
            for pageIndex in range(inReader.getNumPages()):
                outReader.addPage(inReader.getPage(pageIndex))
            outReader.write(outfile)
            outfile.close()
            infile.close()
            return '解密成功，解密后文件为：' + decryptFilename
        else:
            infile.close()
            return filename + "未加密"

    # 加密
    if type_ == 'encrypt':
        # 加密后的文件名
        encryptFilename = os.path.join(path, basename[:len(basename) - 4] + "_encrypt.pdf")
        outfile = open(encryptFilename, 'wb')  # 二进制写
        outWriter = PyPDF2.PdfFileWriter()
        outWriter.encrypt(pwd)  # 加密
        for pageNum in range(inReader.getNumPages()):
            outWriter.addPage(inReader.getPage(pageNum))
        outWriter.write(outfile)
        outfile.close()
        infile.close()
        return '加密成功，加密后文件为：' + encryptFilename


def rotatePDF(filename, param):
    """
    PDF旋转
    :param filename:
    :param param:
    :return:
    """
    if (int(param) % 90) != 0:
        return "旋转失败，旋转度数必须是90的整数倍"

    infile = open(filename, 'rb')  # 二进制读
    inReader = PyPDF2.PdfReader(infile)

    basename = os.path.basename(filename)  # 取文件名
    path = os.path.dirname(filename)  # 取路径
    # 旋转后的文件名
    rotateFilename = os.path.join(path, basename[:len(basename) - 4] + "_rotate.pdf")
    outfile = open(rotateFilename, 'wb')  # 二进制写
    outWriter = PyPDF2.PdfWriter()
    for pageNum in range(len(inReader.pages)):
        page = inReader.pages[pageNum]
        page.rotate(int(param))  # 顺时针旋转
        # page.rotateCounterClockwise(int(param))  # 逆时针旋转
        outWriter.add_page(page)
    outWriter.write(outfile)
    outfile.close()
    infile.close()
    return '旋转成功，旋转后文件为：' + rotateFilename


def watermarkPDF(filename, templateFile):
    """
    叠加页面：给PDF加水印
    :param filename: 需要加水印的PDF
    :param templateFile: 水印模板PDF
    :return:
    """
    # 获取水印page
    templateFile = open(templateFile, 'rb')
    templateReader = PyPDF2.PdfFileReader(templateFile)
    templatePage = templateReader.getPage(0)

    # 加水印后的PDF
    basename = os.path.basename(filename)  # 取文件名
    path = os.path.dirname(filename)  # 取路径
    # 叠加后的文件名
    markedFileName = os.path.join(path, basename[:len(basename) - 4] + "_watermark.pdf")
    outfile = open(markedFileName, 'wb')
    outWriter = PyPDF2.PdfFileWriter()

    # 需要加水印的PDF
    infile = open(filename, 'rb')
    inReader = PyPDF2.PdfFileReader(infile)
    for i in range(inReader.getNumPages()):
        page = inReader.getPage(i)
        page.mergePage(templatePage)  # 叠加
        outWriter.addPage(page)
    outWriter.write(outfile)

    infile.close()
    outfile.close()
    templateFile.close()
    return '添加水印成功'


def pdf2docx(infile, outfile):
    """
    pdf转word
    :param infile:
    :param outfile:
    :return:
    """
    cv = Converter(infile)
    cv.convert(outfile, start=0, end=None)
    cv.close()


if __name__ == '__main__':
    # readPDF('../Python编程快速上手—让繁琐工作自动化.pdf')

    # createAndCopyPDF('pdf/demo1.pdf', 'pdf/demo2.pdf', 'pdf/demo3.pdf')

    # 加密
    # print(encryptOrDecryptPDF('pdf/demo3.pdf', 'encrypt', '123'))
    # 解密
    # print(encryptOrDecryptPDF('pdf/demo3_encrypt.pdf', 'decrypt', '123'))

    # 顺时针旋转90°
    # print(rotatePDF(r'E:\UKey信息.pdf', 180))

    # 叠加页面：给demo1.pdf加水印
    # print(watermarkPDF('pdf/demo1.pdf', 'pdf/watermark.pdf'))

    # pdf2docx(r'F:\Download\aa.pdf', r'F:\Download\a.docx')

    readPDF(r'F:\Download\111.pdf', 4, 6, r'F:\Download\2.pdf')
